'''
作为预测代码，可用于对一下模型的预测
1. resnet_ct_train.py: ct_train

'''

import os
import pandas as pd
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision.models import resnet50

from sklearn.metrics import classification_report, confusion_matrix



BatchSize = 16
save_path = '/home/zmy/pytorch_code/checkpoint/FPM_undersample_7_20210416_101757'
weight_name = 'FPM_undersample_7-95-60.501234690348305-regular.pth'

weights_path = save_path+'/'+weight_name


# 自定义数据集和数据预处理
class MyDataset(Dataset):

    def __init__(self, datalist):
        self.data_info = datalist


    def __len__(self):
        return len(self.data_info)

    def __getitem__(self, item):
        patientID = self.data_info[item][0]
        label = self.data_info[item][1]
        ct_path = self.data_info[item][2]


        ct = np.load(ct_path)

        ct = ct[np.newaxis, :, :]

        return {'image': torch.from_numpy(ct), 'label': torch.tensor(label)}


# 读取文件列表
def read_csv(data_sets):
    sets_path = '/data1/zmy/data2021/origin_data/divide_csv/five/'

    # 读取数据集
    data_features = []

    for set in data_sets:
        train_data = pd.read_csv(sets_path+set)
        for j in range(len(train_data)):

            # 读取文件地址
            patientid = train_data['patientID'][j]

            ct_path = '/data1/zmy/data2021/origin_data/Slice/'+str(patientid)+'/CTSlice/'

            name_list = os.listdir(ct_path)


            for it in name_list:

                one_feature = [patientid, int(train_data['cancer_type'][j])-1, ct_path+it]

                data_features.append(one_feature)


    return data_features





# 创建网络
def resnet():
    # 创建resnet50网络
    net = resnet50(pretrained=False)
    net.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    num_ftrs = net.fc.in_features
    net.fc = nn.Linear(in_features=num_ftrs, out_features=5, bias=True)

    # 打印网络结构和参数量
    # print(net)
    # print("Total number of paramerters in networks is {}  ".format(sum(x.numel() for x in net.parameters())))
    return net
# 结果预测
def predict():
    device = torch.device("cuda:3")

    # 创建网络结构
    net = resnet().to(device)

    # 加载测试数据
    test_data_sets = ['test.csv']
    test_list = read_csv(test_data_sets)
    test_dataset = MyDataset(test_list)
    testloader = DataLoader(test_dataset, batch_size=BatchSize, shuffle=False, num_workers=2)

    # 加载模型权重
    net.load_state_dict(torch.load(weights_path))
    net.eval()

    # 返回的预测结果及对应标签
    y_preds = []
    y_true = []

    with torch.no_grad():
        for n_iter, data in enumerate(testloader):
            print('iteration:{}\ttotal {} iterations'.format(n_iter + 1, len(testloader)))

            images = data['image'].type(torch.FloatTensor).to(device)
            labels = data['label'].to(device)

            outputs = net(images)
            _, preds = outputs.max(1)

            y_preds.extend(preds.tolist())
            y_true.extend(labels.tolist())

            # print(preds.tolist())
            # print(labels.tolist())
    return y_preds, y_true


# 结果评估
def evaluation(preds, labels):

    # 计算评估指标
    target_names = ['1', '2', '3', '4', '5']
    result_statis = classification_report(y_true=labels, y_pred=preds, target_names=target_names)
    print(result_statis)

    # 计算混淆矩阵
    confusion = confusion_matrix(y_true=labels, y_pred=preds)
    print(confusion)


if __name__ == '__main__':

    # 预测结果
    y_preds, y_true = predict()

    print('model: ', weight_name)

    # 结果评估
    evaluation(y_preds, y_true)


